package com.chukun.flink.stream.state.operator;

import com.chukun.flink.stream.state.key.KeyStateBase;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.state.filesystem.FsStateBackendFactory;
import org.apache.flink.runtime.state.storage.FileSystemCheckpointStorage;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Calendar;
import java.util.LinkedList;
import java.util.concurrent.TimeUnit;

/**
 * @author chukun
 * @version 1.0.0
 * @description 实现CheckpointedFunction接口，使得该函数具备Operator状态
 * @createTime 2022年05月16日 18:08:00
 */
public class CheckpointedMapOperator implements MapFunction<Long, String>, CheckpointedFunction {

    private static final Logger logger = LoggerFactory.getLogger(KeyStateBase.StateSource.class);

    // ListState状态的类型结构
    private transient ListState<Long> checkoutpointedState;

    // 集合中存储本该输出到外部系统的数据
    private LinkedList<Long> bufferedElements;

    // 标志状态的重分配模式
    private boolean isUnion;

    // 是否让函数产生异常
    private boolean isError;

    public CheckpointedMapOperator(boolean isUnion, boolean isError) {
        this.isUnion = isUnion;
        this.isError = isError;
        this.bufferedElements = new LinkedList<>();
    }

    @Override
    public String map(Long value) throws Exception {
        int size = bufferedElements.size();
        if (size >= 10) {
            for (int i = 0; i < size - 9; i++) {
                bufferedElements.poll();
            }
        }
        bufferedElements.add(value);
        if (isError) {
            int seconds = Calendar.getInstance().get(Calendar.SECOND);
            if (seconds >= 50 && seconds <= 51) {
                int i = 1 / 0;
            }
        }
        logger.info("{} map data :{}", Thread.currentThread().getName(), bufferedElements);
        return "集合中第一个元素是:" + bufferedElements.getFirst() +
                "集合中最后一个元素是:" + bufferedElements.getLast() +
                " length is :" + bufferedElements.size();
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {

        logger.info(Thread.currentThread().getName()+":"+context.getCheckpointId() +":"+context.getCheckpointTimestamp()+".....snapshotState");
        logger.info("{} 快照编号 {} 的元素为 : {}", Thread.currentThread().getName(), context.getCheckpointId(), bufferedElements);

        // 清除状态数据
        checkoutpointedState.clear();
        // 将bufferedElements集合的快照存入状态
        checkoutpointedState.addAll(bufferedElements);

    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {

        // 创建一个ListStateDescriptor状态描述符，指定名称为CheckpointedMapOperator-ListState，类型为 Long
        ListStateDescriptor<Long> descriptor = new ListStateDescriptor<>(
                "CheckpointedMapOperator-ListState",
                TypeInformation.of(new TypeHint<Long>() {}));
        if (isUnion) {
            // 以Union的方式获取状态
            checkoutpointedState = context.getOperatorStateStore().getUnionListState(descriptor);
        } else {
            // 以Even-Split方式获取状态
            checkoutpointedState = context.getOperatorStateStore().getListState(descriptor);
        }
        // 判断作业是否失败后重新恢复
        if (context.isRestored()) {
            logger.info("{} operator状态恢复", Thread.currentThread().getName());
            for (Long element : checkoutpointedState.get()) {
                // 获取状态中的数据，添加进bufferedElements集合中
                bufferedElements.add(element);
            }
        }
        logger.info("{} operator 初始化/恢复{}", Thread.currentThread().getName(), bufferedElements);
    }

    /**
     * 定义数据源
     */
    public static class CustomSource extends RichSourceFunction<Long> {

        public static final Logger logger = LoggerFactory.getLogger(CustomSource.class);

        private static final long serialVersionUID = 1L;

        /**
         * @date: 2020/10/15 18:33
         */
        @Override
        public void open(Configuration parameters) throws Exception {
            TimeUnit.SECONDS.sleep(10);
        }

        /**
         * @date: 2020/10/15 18:33
         */
        @Override
        public void run(SourceContext<Long> ctx) throws Exception {
            Long offset = 0L;
            while (true) {
                logger.info("{}{}{}", Thread.currentThread().getName(), ":发送数据:", offset);
                ctx.collect(offset);
                offset += 1;
                TimeUnit.SECONDS.sleep(1);
            }
        }

        /**
         */
        @Override
        public void cancel() {}
    }

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // 设置每10秒执行一次检查点操作
        env.enableCheckpointing(10000);

        // 设置作业的全局并行度为2
        env.setParallelism(2);

        String path = "file:///E:/checkoutpoint/CheckpointedMapOperator";
        // 设置FileSystemCheckpointStorage作为状态后端，检查点文件存储在指定的路径
        env.getCheckpointConfig().setCheckpointStorage(new FileSystemCheckpointStorage(path));

        // 设置作业失败的重启策略
        env.setRestartStrategy(RestartStrategies.failureRateRestart(3, Time.of(5, TimeUnit.MINUTES),Time.of(10, TimeUnit.SECONDS)));

        DataStream<Long> streamSource = env.addSource(new CustomSource()).setParallelism(1);

        DataStream<String> mapResult = streamSource.map(new CheckpointedMapOperator(false, true));

        mapResult.print("输出结果");

        env.execute("CheckpointedMapOperator");

    }
}
